import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
import mindspore.dataset as ds
import numpy as np


def create_data(data, w=2.0, b=3.0):
    for item in range(data):
        # 使用uniform均匀分布数据
        x = np.random.uniform(-10.0, 10.0)
        # 指定一个噪声来模拟数据
        # 使用normal使其满足高斯分布,参数为数据的中间值和浮动值
        noise = np.random.normal(0, 1)
        y = w * x + b + noise
        # 使用yield返回数据生成器，方便后续对数据进行处理
        yield np.array([x]).astype(np.float32), np.array([y]).astype(np.float32)


# 处理数据集
def create_dataset(data_num, batch_size, repeat_num):
    dataset = ds.GeneratorDataset(list(create_data(data_num)), column_names=['data', 'label'], shuffle=True)
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.repeat(repeat_num)
    return dataset


# 创建训练数据集
train_dataset = create_dataset(data_num=1000, batch_size=10, repeat_num=1)


class MyNet(nn.Cell):
    """定义网络模型"""
    def __init__(self):
        super(MyNet, self).__init__()
        self.layer = nn.Dense(1, 1)

    def construct(self, x):
        fx = self.layer(x)
        return fx


# 初始化网络模型
net = MyNet()

# 获取可训练参数w和b
model_params = net.trainable_params()


class MyLossFunction(nn.LossBase):
    """定义损失函数"""
    def __init__(self):
        super(MyLossFunction, self).__init__()
        self.abs = ops.Abs()

    def construct(self, target, predict):
        x = self.abs(target - predict)
        return self.get_loss(x) # 计算加权损失


class MyOptimizer(nn.Optimizer):
    """自定义优化器"""
    # 我的建议是直接nn.Momentum
    def __init__(self, learning_rate, params, momentum=0.9):
        super(MyOptimizer, self).__init__(learning_rate, params)
        self.assign = ops.Assign()
        # 定义初速度
        self.moment = ms.Parameter(ms.Tensor(momentum, ms.float32), name="moment")
        # 定义动量参数
        self.momentum = self.parameters.clone(prefix="momentum", init="zeros")

    def construct(self, gradients):
        lr = self.get_lr()
        params = self.parameters
        for i in range(len(params)):
            # 给网络赋值动量参数
            self.assign(self.momentum[i], self.momentum[i] * self.moment + gradients[i])
            # 计算新的权重参数
            update = params[i] - self.momentum[i] * lr
            # 更新权重参数
            self.assign(params[i], update)
        return params


class MyLossNet(nn.Cell):
    """定义损失网络"""
    def __init__(self, backbone, loss_fn): # backbone翻译为主干网络
        super(MyLossNet, self).__init__(auto_prefix=False)
        self.backbone = backbone
        self.loss_fn = loss_fn

    def construct(self, data, label):
        # 将神经网络的输出输入到损失函数，返回损失值
        out = self.backbone(data)
        return self.loss_fn(out, label)

    def backbone_net(self):
        return self.backbone


class MyTrainStep(nn.TrainOneStepCell):
    """定义模型训练过程"""
    def __init__(self, network, optimizer):
        # 初始化参数
        super(MyTrainStep, self).__init__(network, optimizer)
        self.grad = ops.GradOperation(get_by_list=True)

    def construct(self, data, label):
        # 构建训练过程
        weights = self.weights
        loss = self.network(data, label)
        grads = self.grad(self.network, weights)(data, label)
        return loss, self.optimizer(grads)


# 执行训练
loss_func = MyLossFunction()
optim = MyOptimizer(0.01, net.trainable_params())
loss_net = MyLossNet(net, loss_func)
train_net = MyTrainStep(loss_net, optim)

for data in train_dataset.create_dict_iterator():
    train_net(data['data'], data['label'])
    loss = loss_net(data['data'], data['label'])
    print(f"Loss: {loss}")


class MyMAE(nn.Metric):
    """定义评价指标"""
    def __init__(self):
        super(MyMAE, self).__init__()
        self.clear()

    def clear(self):
        # 初始化参数
        self.abs_error_sum = 0
        self.samples_num = 0

    def update(self, *inputs):
        # 更新参数
        y_predict = inputs[0].asnumpy()
        y = inputs[1].asnumpy()

        # 计算绝对误差
        error_abs = np.abs(y.reshape(y_predict.shape) - y_predict)
        self.abs_error_sum += error_abs.sum()
        self.samples_num += y.shape[0]

    def eval(self):
        # 计算最终评估结果
        return self.abs_error_sum / self.samples_num


class MyEvalNet(nn.Cell):
    """定义评估网络"""
    def __init__(self, network):
        super(MyEvalNet, self).__init__()
        self.network = network

    def construct(self, data, label):
        outputs = self.network(data)
        return outputs, label


# 加载数据集
eval_dataset = create_dataset(data_num=160, batch_size=16, repeat_num=1)

# 初始化网络
eval_net = MyEvalNet(net)
eval_net.set_train(False)
mae = MyMAE()

# 执行推理
for data in eval_dataset.create_dict_iterator():
    output, y_eval = eval_net(data['data'], data['label'])
    mae.update(output, y_eval)

acc = mae.eval()
print("MAE:", acc)


ms.save_checkpoint(net, "./linear.ckpt")
param_dict = ms.load_checkpoint("./linear.ckpt")

for i in param_dict:
    print(i, ":", param_dict[i].asnumpy())


